package cn.edu.lingnan.utils;

import okhttp3.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class OpenAIChatStreaming {
    private static final String API_URL = "https://openrouter.ai/api/v1/chat/completions";
    private static final String API_KEY = "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";
    private static final String MODEL = "deepseek/deepseek-chat-v3-0324:free";
    private static final OkHttpClient HTTP_CLIENT = new OkHttpClient();
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper()
            .disable(SerializationFeature.FAIL_ON_EMPTY_BEANS);

    public static String chatStream(String question) throws IOException {
        // 1. 构造消息列表
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage("user", question));

        // 2. 构建请求JSON
        String requestBody = OBJECT_MAPPER.writeValueAsString(
                new ChatRequest(MODEL, messages, true)
        );

        // 3. 发送请求
        Request request = new Request.Builder()
                .url(API_URL)
                .addHeader("Authorization", "Bearer " + API_KEY)
                .addHeader("Content-Type", "application/json")
                .post(RequestBody.create(requestBody, MediaType.parse("application/json")))
                .build();

        // 4. 处理流式响应
        try (Response response = HTTP_CLIENT.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("请求失败: " + response.code() + " - " + response.message());
            }

            ResponseBody body = response.body();
            if (body == null) throw new IOException("响应体为空");

            StringBuilder responseBuilder = new StringBuilder();

            try (var source = body.source()) {
                while (!source.exhausted()) {
                    String line = source.readUtf8Line();
                    if (line == null || line.isEmpty() || line.equals("data: [DONE]")) continue;

                    if (line.startsWith("data: ")) {
                        String jsonData = line.substring(6).trim();
                        if (jsonData.isEmpty()) continue;

                        try {
                            var node = OBJECT_MAPPER.readTree(jsonData);
                            String content = node.path("choices").get(0).path("delta").path("content").asText();
                            if (content != null && !content.isEmpty()) {
                                responseBuilder.append(content);
                            }
                        } catch (Exception e) {
                            System.err.println("[解析错误] " + e.getMessage());
                        }
                    }
                }
            }
            return responseBuilder.toString();
        }
    }

    // ===== 内部类保持不变 =====
    private static class ChatMessage {
        public String role;
        public String content;

        public ChatMessage(String role, String content) {
            this.role = role;
            this.content = content;
        }
    }

    private static class ChatRequest {
        public String model;
        public List<ChatMessage> messages;
        public boolean stream;

        public ChatRequest(String model, List<ChatMessage> messages, boolean stream) {
            this.model = model;
            this.messages = messages;
            this.stream = stream;
        }
    }

    public static void main(String[] args) {
        try {
            // 修改后的调用方式
            //计算请求开始时间
            long startTime = System.currentTimeMillis();
            String response = OpenAIChatStreaming.chatStream("你好");
            System.out.println("完整响应：\n" + response);
            long endTime = System.currentTimeMillis();
            System.out.println("请求耗时：" + (endTime - startTime) + "毫秒");
        } catch (IOException e) {
            System.err.println("发生错误: " + e.getMessage());
        }
    }
}